"""
python fujian3_cal_average_sales_gragh.py
"""

import os
import json
import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import uniform_filter1d
from tqdm import tqdm

# 输入和输出路径
input_directory = '../fujian/fujian2/RandomForest'
output_file_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
plot_directory = os.path.join(os.path.dirname(output_file_path), 'sales_gragh')

# 确保输出文件夹和图表文件夹存在
os.makedirs(plot_directory, exist_ok=True)

# 存储每个 category 的均值
average_sales = []

# 遍历输入目录下的所有 JSON 文件并显示进度条
for filename in tqdm(os.listdir(input_directory), desc="Processing categories"):
    # 使用正则表达式提取 category 编号
    match = re.search(r'category(\d+)', filename)
    if match:
        category_id = match.group(1)  # 获取类别编号
        
        # 构建文件路径
        file_path = os.path.join(input_directory, filename)

        # 读取 JSON 文件
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        
        # 提取销售数据并转换为 DataFrame
        dates = [item['date'] for item in data]
        sales = np.array([item['predicted_sales'] for item in data])
        df = pd.DataFrame({'date': dates, 'sales': sales})
        
        # 计算移动平均和指数平滑数据
        window_size = 5  # 可调整窗口大小
        moving_avg = uniform_filter1d(sales, size=window_size)
        exp_smooth = pd.Series(sales).ewm(span=window_size, adjust=False).mean().values
        
        # 计算处理后数据的均值
        average_sales_value = np.mean(exp_smooth)
        
        # 添加结果到列表
        average_sales.append({
            "category_id": category_id,
            "average_sales": average_sales_value
        })
        
        # 绘制图表
        plt.figure(figsize=(10, 6))
        
        # 原始数据的散点图
        plt.scatter(df['date'], df['sales'], color='blue', label='Original Data', alpha=0.6)
        
        # 经过移动平均和指数平滑的数据的散点图
        plt.plot(df['date'], moving_avg, color='orange', linestyle='-', label='Moving Average')
        plt.plot(df['date'], exp_smooth, color='green', linestyle='-', label='Exponential Smoothing')
        
        # 绘制均值水平线
        plt.axhline(average_sales_value, color='red', linestyle='--', label=f'Average: {average_sales_value:.2f}')
        
        # 图表格式设置
        plt.title(f'Category {category_id} Sales Data')
        plt.xlabel('Date')
        plt.ylabel('Sales')
        plt.xticks(rotation=45)
        plt.legend()
        
        # 保存图表
        plot_path = os.path.join(plot_directory, f'sales_plot_category_{category_id}.png')
        plt.savefig(plot_path, bbox_inches='tight')
        plt.close()
        
# 输出均值到 JSON 文件
with open(output_file_path, 'w', encoding='utf-8') as output_file:
    json.dump(average_sales, output_file, indent=4)

print("Average sales calculated, plots generated, and data saved to", output_file_path)
